{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Layers\n",
    "\n",
    "> Basic layers for constructor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "# from nbdev.showdoc import *\n",
    "from fastcore.test import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "from torch.nn.utils import spectral_norm # weight_norm, \n",
    "from collections import OrderedDict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ConvLayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "act_fn = nn.ReLU(inplace=True)\n",
    "\n",
    "class ConvLayer(nn.Sequential):\n",
    "    \"\"\"Basic conv layers block\"\"\"\n",
    "    def __init__(self, ni, nf, ks=3, stride=1, \n",
    "            act=True,  act_fn=act_fn, \n",
    "            bn_layer=True, bn_1st=True, zero_bn=False, \n",
    "            padding=None, bias=False, groups=1, **kwargs):\n",
    "\n",
    "#         self.act = act\n",
    "        if padding==None: padding = ks//2  \n",
    "        layers = [('conv', nn.Conv2d(ni, nf, ks, stride=stride, padding=padding, bias=bias, groups=groups))]\n",
    "        act_bn = [('act_fn', act_fn)] if act else []\n",
    "        if bn_layer:\n",
    "            bn = nn.BatchNorm2d(nf)\n",
    "            nn.init.constant_(bn.weight, 0. if zero_bn else 1.) \n",
    "            act_bn += [('bn', bn)]\n",
    "        if bn_1st: act_bn.reverse()\n",
    "        layers += act_bn\n",
    "        super().__init__(OrderedDict(layers))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConvLayer(\n",
       "  (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "  (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (act_fn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv_layer = ConvLayer(32, 64)\n",
    "conv_layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConvLayer(\n",
       "  (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "  (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv_layer = ConvLayer(32, 64, act=False)\n",
    "conv_layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConvLayer(\n",
       "  (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "  (act_fn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv_layer = ConvLayer(32, 64, bn_layer=False)\n",
    "conv_layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConvLayer(\n",
       "  (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "  (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (act_fn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv_layer = ConvLayer(32, 64, bn_1st=True)\n",
    "conv_layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConvLayer(\n",
       "  (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "  (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (act_fn): LeakyReLU(negative_slope=0.01)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv_layer = ConvLayer(32, 64, bn_1st=True, act_fn=nn.LeakyReLU())\n",
    "conv_layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Flatten"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class Flatten(nn.Module):\n",
    "    '''flat x to vector'''\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "    def forward(self, x): return x.view(x.size(0), -1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Noop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "def noop(x): return x\n",
    "\n",
    "class Noop(nn.Module): \n",
    "    '''Dummy module'''\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Downsample block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def DownsampleBlock(conv_layer, ni, nf, ks, stride, act=False, **kwargs):\n",
    "    '''Base downsample for res-like blocks'''\n",
    "    return conv_layer(ni, nf, ks, stride, act, **kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BasicBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class BasicBlock(nn.Module):\n",
    "    \"\"\"Basic block (simplified) as in pytorch resnet\"\"\"\n",
    "    def __init__(self, ni, nf,  expansion=1, stride=1, zero_bn=False,\n",
    "                conv_layer=ConvLayer, act_fn=act_fn, \n",
    "                downsample_block=DownsampleBlock, **kwargs):\n",
    "        super().__init__()\n",
    "        self.downsample = not ni==nf or stride==2\n",
    "        self.conv = nn.Sequential(OrderedDict([\n",
    "            ('conv_0', conv_layer(ni, nf, stride=stride, act_fn=act_fn,  **kwargs)),\n",
    "            ('conv_1', conv_layer(nf, nf, zero_bn=zero_bn, act=False, act_fn=act_fn, **kwargs))]))\n",
    "        if self.downsample:\n",
    "            self.downsample = downsample_block(conv_layer, ni, nf, ks=1, stride=stride, act=False, **kwargs)\n",
    "        self.merge = Noop()\n",
    "        self.act_conn = act_fn\n",
    "                         \n",
    "    def forward(self, x):\n",
    "        identity = x\n",
    "        out = self.conv(x)\n",
    "        if self.downsample:\n",
    "            identity = self.downsample(x)\n",
    "        return self.act_conn(self.merge(out + identity))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BasicBlock(\n",
       "  (conv): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (merge): Noop()\n",
       "  (act_conn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b_block = BasicBlock(64,64)\n",
    "b_block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 64, 32, 32])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = b_block(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BasicBlock(\n",
       "  (conv): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (downsample): ConvLayer(\n",
       "    (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "    (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (merge): Noop()\n",
       "  (act_conn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b_block = BasicBlock(64,64, stride=2)\n",
    "b_block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 64, 16, 16])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = b_block(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BasicBlock(\n",
       "  (conv): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (downsample): ConvLayer(\n",
       "    (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "    (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (merge): Noop()\n",
       "  (act_conn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b_block = BasicBlock(64,128, stride=2)\n",
    "b_block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 128, 16, 16])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = b_block(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BasicBlock(\n",
       "  (conv): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (downsample): ConvLayer(\n",
       "    (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "    (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (merge): Noop()\n",
       "  (act_conn): LeakyReLU(negative_slope=0.01)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b_block = BasicBlock(64,128, stride=2, act_fn=nn.LeakyReLU())\n",
    "b_block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 128, 16, 16])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = b_block(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BasicBlock(\n",
       "  (conv): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (downsample): ConvLayer(\n",
       "    (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "    (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (merge): Noop()\n",
       "  (act_conn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b_block = BasicBlock(64,128, stride=2, bn_1st=False)\n",
    "b_block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 128, 16, 16])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = b_block(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bottleneck"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class Bottleneck(nn.Module):\n",
    "    '''Bottlneck block for resnet models'''\n",
    "    def __init__(self, ni, nh, expansion=4, stride=1, zero_bn=False, \n",
    "                conv_layer=ConvLayer, act_fn=act_fn,\n",
    "                 downsample_block=DownsampleBlock, **kwargs):\n",
    "#                  groups=1, base_width=64, dilation=1, norm_layer=None\n",
    "        super().__init__()\n",
    "        self.downsample = not ni==nh or stride==2\n",
    "        ni = ni*expansion\n",
    "        nf = nh*expansion\n",
    "        self.conv = nn.Sequential(OrderedDict([\n",
    "            ('conv_0', conv_layer(ni, nh, ks=1,            act_fn=act_fn, **kwargs)),\n",
    "            ('conv_1', conv_layer(nh, nh, stride=stride,   act_fn=act_fn, **kwargs)),\n",
    "            ('conv_2', conv_layer(nh, nf, ks=1, zero_bn=zero_bn, act=False, act_fn=act_fn, **kwargs))]))\n",
    "        if self.downsample:\n",
    "            self.downsample = downsample_block(conv_layer, ni, nf, ks=1, stride=stride, act=False, act_fn=act_fn, **kwargs)\n",
    "        self.merge = Noop()\n",
    "        self.act_conn = act_fn\n",
    "\n",
    "    def forward(self, x):\n",
    "        identity = x\n",
    "        out = self.conv(x)\n",
    "        if self.downsample:\n",
    "            identity = self.downsample(x)\n",
    "        return self.act_conn(self.merge(out + identity))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Bottleneck(\n",
       "  (conv): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_2): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (downsample): ConvLayer(\n",
       "    (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (merge): Noop()\n",
       "  (act_conn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b_block = Bottleneck(16,64)\n",
    "b_block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 256, 32, 32])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = b_block(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Bottleneck(\n",
       "  (conv): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_2): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (merge): Noop()\n",
       "  (act_conn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b_block = Bottleneck(64,64)\n",
    "b_block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 256, 32, 32])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 256, 32, 32)\n",
    "y = b_block(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Bottleneck(\n",
       "  (conv): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (conv_2): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (merge): Noop()\n",
       "  (act_conn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b_block = Bottleneck(64,64, bn_1st=False)\n",
    "b_block"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# XResBlock"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As in fastai v1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DownsampleLayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class DownsampleLayer(nn.Sequential):\n",
    "    \"\"\"Downsample layer for Xresnet Resblock\"\"\"\n",
    "    def __init__(self, conv_layer, ni, nf, stride, act,  \n",
    "                 pool=nn.AvgPool2d(2, ceil_mode=True), pool_1st=True,\n",
    "                 **kwargs):\n",
    "        layers  = [] if stride==1 else [('pool', pool)]\n",
    "        layers += [] if ni==nf else [('idconv', conv_layer(ni, nf, 1, act=act, **kwargs))]\n",
    "        if not pool_1st: layers.reverse()\n",
    "        super().__init__(OrderedDict(layers))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## XresBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class XResBlock(nn.Module):\n",
    "    def __init__(self, ni, nh, expansion=1, stride=1, zero_bn=True, \n",
    "                 conv_layer=ConvLayer, act_fn=act_fn, **kwargs):\n",
    "        super().__init__()\n",
    "        nf,ni = nh*expansion,ni*expansion\n",
    "        layers  = [('conv_0', conv_layer(ni, nh, 3, stride=stride, act_fn=act_fn, **kwargs)),\n",
    "                   ('conv_1', conv_layer(nh, nf, 3, zero_bn=zero_bn, act=False, act_fn=act_fn, **kwargs))\n",
    "        ] if expansion == 1 else [\n",
    "                   ('conv_0', conv_layer(ni, nh, 1, act_fn=act_fn, **kwargs)),\n",
    "                   ('conv_1', conv_layer(nh, nh, 3, stride=stride, act_fn=act_fn, **kwargs)),\n",
    "                   ('conv_2', conv_layer(nh, nf, 1, zero_bn=zero_bn, act=False, act_fn=act_fn, **kwargs))\n",
    "        ]\n",
    "        self.convs = nn.Sequential(OrderedDict(layers))\n",
    "        self.identity = DownsampleLayer(conv_layer, ni, nf, stride, act=False, act_fn=act_fn, **kwargs) if ni!=nf or stride==2 else Noop()       \n",
    "        self.merge = Noop() # us it to visualize in repr residual connection\n",
    "        self.act_fn = act_fn\n",
    "\n",
    "    def forward(self, x): return self.act_fn(self.merge(self.convs(x) + self.identity(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_block = XResBlock(16,64,expansion=4, stride=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "XResBlock(\n",
       "  (convs): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_2): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (identity): DownsampleLayer(\n",
       "    (pool): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
       "    (idconv): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (merge): Noop()\n",
       "  (act_fn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res_block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 256, 16, 16])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = res_block(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert y.shape==torch.Size([bs, 256, 16, 16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_block = XResBlock(16,64,expansion=4, stride=2, pool_1st=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "XResBlock(\n",
       "  (convs): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_2): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (identity): DownsampleLayer(\n",
       "    (idconv): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (pool): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
       "  )\n",
       "  (merge): Noop()\n",
       "  (act_fn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res_block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 256, 16, 16])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = res_block(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert y.shape==torch.Size([bs, 256, 16, 16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_block = XResBlock(16,64,expansion=4, stride=2, bn_1st=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "XResBlock(\n",
       "  (convs): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (conv_2): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (identity): DownsampleLayer(\n",
       "    (pool): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
       "    (idconv): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (merge): Noop()\n",
       "  (act_fn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res_block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 256, 16, 16])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = res_block(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert y.shape==torch.Size([bs, 256, 16, 16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_block = XResBlock(16,64,expansion=4, stride=2, pool=nn.MaxPool2d(kernel_size=3, stride=2, padding=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "XResBlock(\n",
       "  (convs): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_2): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (identity): DownsampleLayer(\n",
       "    (pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "    (idconv): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (merge): Noop()\n",
       "  (act_fn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res_block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 256, 16, 16])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = res_block(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert y.shape==torch.Size([bs, 256, 16, 16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_block = XResBlock(16,64,expansion=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "XResBlock(\n",
       "  (convs): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_2): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (identity): DownsampleLayer(\n",
       "    (idconv): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (merge): Noop()\n",
       "  (act_fn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res_block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 256, 32, 32])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = res_block(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert y.shape==torch.Size([bs, 256, 32, 32])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_block = XResBlock(64,64,expansion=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 64, 32, 32])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = res_block(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert y.shape==torch.Size([bs, 64, 32, 32])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_block = XResBlock(64,64,expansion=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 256, 32, 32])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 256, 32, 32)\n",
    "y = res_block(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert y.shape==torch.Size([bs, 256, 32, 32])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SimpleSelfAttention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "# SA module from mxresnet at fastai. todo - add persons!!!\n",
    "\n",
    "#Unmodified from https://github.com/fastai/fastai/blob/5c51f9eabf76853a89a9bc5741804d2ed4407e49/fastai/layers.py\n",
    "def conv1d(ni:int, no:int, ks:int=1, stride:int=1, padding:int=0, bias:bool=False):\n",
    "    \"Create and initialize a `nn.Conv1d` layer with spectral normalization.\"\n",
    "    conv = nn.Conv1d(ni, no, ks, stride=stride, padding=padding, bias=bias)\n",
    "    nn.init.kaiming_normal_(conv.weight)\n",
    "    if bias: conv.bias.data.zero_()\n",
    "    return spectral_norm(conv)\n",
    "\n",
    "\n",
    "# Adapted from SelfAttention layer at https://github.com/fastai/fastai/blob/5c51f9eabf76853a89a9bc5741804d2ed4407e49/fastai/layers.py\n",
    "# Inspired by https://arxiv.org/pdf/1805.08318.pdf\n",
    "class SimpleSelfAttention(nn.Module):\n",
    "    def __init__(self, n_in:int, ks=1, sym=False):#, n_out:int):\n",
    "        super().__init__()\n",
    "        self.conv = conv1d(n_in, n_in, ks, padding=ks//2, bias=False)      \n",
    "        self.gamma = nn.Parameter(torch.tensor([0.]))\n",
    "        self.sym = sym\n",
    "        self.n_in = n_in\n",
    "    def forward(self,x):\n",
    "        if self.sym:\n",
    "            # symmetry hack by https://github.com/mgrankin\n",
    "            c = self.conv.weight.view(self.n_in,self.n_in)\n",
    "            c = (c + c.t())/2\n",
    "            self.conv.weight = c.view(self.n_in,self.n_in,1)\n",
    "        size = x.size()  \n",
    "        x = x.view(*size[:2],-1)   # (C,N)\n",
    "        # changed the order of mutiplication to avoid O(N^2) complexity\n",
    "        # (x*xT)*(W*x) instead of (x*(xT*(W*x)))\n",
    "        convx = self.conv(x)   # (C,C) * (C,N) = (C,N)   => O(NC^2)\n",
    "        xxT = torch.bmm(x,x.permute(0,2,1).contiguous())   # (C,N) * (N,C) = (C,C)   => O(NC^2)\n",
    "        o = torch.bmm(xxT, convx)   # (C,C) * (C,N) = (C,N)   => O(NC^2)\n",
    "        o = self.gamma * o + x\n",
    "        return o.view(*size).contiguous()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convolutions blocks for ResBlock."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class ConvBlockBasic(nn.Sequential):\n",
    "    '''Block of conv layers for basic resblock'''\n",
    "    def __init__(self, conv_layer,ni,nh,nf,stride,zero_bn,act=False,**kwargs):\n",
    "        super().__init__(OrderedDict([\n",
    "            ('conv_0', conv_layer(ni, nh, 3, stride=stride, **kwargs)),\n",
    "            ('conv_1', conv_layer(nh, nf, 3, zero_bn=zero_bn, act=act, **kwargs))\n",
    "                                ]))\n",
    "\n",
    "class ConvBlockBottle(nn.Sequential):\n",
    "    '''Basic block of conv layers for resblock'''\n",
    "    def __init__(self, conv_layer,ni,nh,nf,stride,zero_bn,act=False,**kwargs):\n",
    "        super().__init__(OrderedDict([\n",
    "            ('conv_0', conv_layer(ni, nh, 1,                        **kwargs)),\n",
    "            ('conv_1', conv_layer(nh, nh, 3, stride=stride,         **kwargs)),\n",
    "            ('conv_2', conv_layer(nh, nf, 1, zero_bn=zero_bn, act=act, **kwargs))\n",
    "                                ]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConvBlockBasic(\n",
       "  (conv_0): ConvLayer(\n",
       "    (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (act_fn): ReLU(inplace=True)\n",
       "  )\n",
       "  (conv_1): ConvLayer(\n",
       "    (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ConvBlockBasic(ConvLayer,64,64,64,1,True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ResBlock."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class ResBlock(nn.Module):\n",
    "    def __init__(self, ni, nh, expansion=1, stride=1, conv_layer=ConvLayer,\n",
    "                 conv_block=ConvBlockBasic, downsample=DownsampleLayer,\n",
    "                 act_fn=act_fn, act_id=False, zero_bn=True, sa=False, sym=False, **kwargs):\n",
    "        super().__init__()\n",
    "        nf,ni = nh*expansion,ni*expansion\n",
    "        self.convs = conv_block(conv_layer,ni,nh,nf,stride,zero_bn,act=act_id,act_fn=act_fn,**kwargs)\n",
    "        self.sa = SimpleSelfAttention(nf,ks=1,sym=sym) if sa else noop\n",
    "        self.identity = downsample(conv_layer,ni,nf,stride,act=act_id,act_fn=act_fn,**kwargs) if ni!=nf or stride==2 else Noop()       \n",
    "        self.merge =  act_fn if not act_id else Noop()\n",
    "    def forward(self, x): return self.merge(self.sa(self.convs(x)) + self.identity(x)) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResBlock(\n",
       "  (convs): ConvBlockBasic(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (sa): SimpleSelfAttention(\n",
       "    (conv): Conv1d(64, 64, kernel_size=(1,), stride=(1,), bias=False)\n",
       "  )\n",
       "  (identity): DownsampleLayer(\n",
       "    (pool): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
       "  )\n",
       "  (merge): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "basicblock = ResBlock(64,64,1,stride=2, sa=True)\n",
    "basicblock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([8, 64, 16, 16])\n"
     ]
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = basicblock(xb)\n",
    "print(y.shape)\n",
    "assert y.shape==torch.Size([bs, 64, 16, 16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResBlock(\n",
       "  (convs): ConvBlockBasic(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (sa): SimpleSelfAttention(\n",
       "    (conv): Conv1d(128, 128, kernel_size=(1,), stride=(1,), bias=False)\n",
       "  )\n",
       "  (identity): DownsampleLayer(\n",
       "    (pool): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
       "    (idconv): ConvLayer(\n",
       "      (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (merge): Noop()\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "basicblock = ResBlock(64,128,1,stride=2, sa=True, act_id=True)\n",
    "basicblock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([8, 128, 16, 16])\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = basicblock(xb)\n",
    "print(y.shape)\n",
    "assert y.shape==torch.Size([bs, 128, 16, 16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResBlock(\n",
       "  (convs): ConvBlockBasic(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (identity): DownsampleLayer(\n",
       "    (pool): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
       "    (idconv): ConvLayer(\n",
       "      (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (merge): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "basicblock = ResBlock(64,128,1,stride=2)\n",
    "basicblock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([8, 128, 16, 16])\n"
     ]
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = basicblock(xb)\n",
    "print(y.shape)\n",
    "assert y.shape==torch.Size([bs, 128, 16, 16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "act_fn = nn.LeakyReLU(inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResBlock(\n",
       "  (convs): ConvBlockBottle(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    )\n",
       "    (conv_2): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (identity): DownsampleLayer(\n",
       "    (idconv): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (merge): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bottle = ResBlock(16, 64, 4, 1,conv_block=ConvBlockBottle, act_fn=act_fn)\n",
    "bottle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 128, 16, 16])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 8\n",
    "xb = torch.randn(bs, 64, 32, 32)\n",
    "y = basicblock(xb)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert y.shape==torch.Size([bs, 128, 16, 16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResBlock(\n",
       "  (convs): ConvBlockBasic(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (sa): SimpleSelfAttention(\n",
       "    (conv): Conv1d(128, 128, kernel_size=(1,), stride=(1,), bias=False)\n",
       "  )\n",
       "  (identity): DownsampleLayer(\n",
       "    (pool): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
       "    (idconv): ConvLayer(\n",
       "      (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (merge): Noop()\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "basicblock = ResBlock(64,128,1,stride=2, sa=True, act_id=True, act_fn=act_fn)\n",
    "basicblock"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# end\n",
    "model_constructor\n",
    "by ayasyrev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converted 00_constructor.ipynb.\n",
      "Converted 01_layers.ipynb.\n",
      "Converted 02_resnet.ipynb.\n",
      "Converted 03_xresnet.ipynb.\n",
      "Converted 80_test_layers.ipynb.\n",
      "Converted 81_test_xresnet.ipynb.\n",
      "Converted index.ipynb.\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "from nbdev.export import *\n",
    "notebook2script()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
